In [1]:
# import packages
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from matplotlib.ticker import AutoMinorLocator
from matplotlib import rcParams
rcParams['ytick.direction'] = 'out'
from positronium import Bohr
In [2]:
def n_plot(ax, n_min = 1, n_max=100, **kwargs):
en = Bohr.energy(float('inf'), np.arange(n_min, n_max), unit='eV')
ax.hlines(en, 0, 1, **kwargs)
return ax
def n_labels(ax, n_vals, x_pos, **kwargs):
labels = ['$n=\infty$' if n == float('inf') else '$n=$%d'%n for n in n_vals]
for n, lbl in zip(n_vals, labels):
en = Bohr.energy(float('inf'), n, unit='eV')
ax.annotate(lbl, (x_pos, en), **kwargs)
return ax
In [3]:
# plot setup
fig, ax = plt.subplots(figsize=(4, 6))
n_plot(ax, 1, 400, lw=2, color='black', alpha=0.3)
n_labels(ax, [1, 2, 3, 4, float('inf')], 1.05, va='center', color='black')
# format
ax.set_ylabel("energy (eV)")
ax.set_xlim(-0.1, 1.2)
ax.set_ylim(-7, 0)
ax.set_xticks([])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
# output
plt.tight_layout()
plt.show()
In [ ]: